# -*- coding: utf-8 -*-

import tensorflow as tf
import h5py
import numpy as np


def main():
    weights_name = ["conv1_1/weights", "conv1_2/weights",
                    "conv2_1/weights", "conv2_2/weights",
                    "conv3_1/weights", "conv3_2/weights", "conv3_3/weights",
                    "conv4_1/weights", "conv4_2/weights", "conv4_3/weights",
                    "conv5_1/weights", "conv5_2/weights", "conv5_3/weights", "conv5_4/weights"]
    biases_name = ["conv1_1/biases", "conv1_2/biases",
                   "conv2_1/biases", "conv2_2/biases",
                   "conv3_1/biases", "conv3_2/biases", "conv3_3/biases",
                   "conv4_1/biases", "conv4_2/biases", "conv4_3/biases",
                   "conv5_1/biases", "conv5_2/biases", "conv5_3/biases"]
    check_point_file_path = "./model/vehicle_L2Loss.npk-2700"
    reader = tf.train.NewCheckpointReader(check_point_file_path)
    # hdf5_file = h5py.File("vehicle_L2Loss.h5", "w")
    # for weight_name in weights_name:
    #     hdf5_file.create_dataset(name=weight_name, data=reader.get_tensor(weight_name))
    # for bias_name in biases_name:
    #     hdf5_file.create_dataset(name=bias_name, data=reader.get_tensor(bias_name))
    # hdf5_file.close()

    hdf5_file = h5py.File("./model/vehicle_L2Loss.h5", "r")
    weights_src = reader.get_tensor("conv4_3/weights")
    biases_src = reader.get_tensor("conv4_3/biases")
    weights_h5 = hdf5_file["conv4_3/weights"]
    biases_h5 = hdf5_file["conv4_3/biases"]
    print(weights_src.shape, weights_h5.shape)
    print(biases_src.shape, biases_h5.shape)
    print(np.sum(weights_src - weights_h5))
    print(np.sum(biases_src - biases_h5))
    print(np.array(biases_h5))
    hdf5_file.close()


if __name__ == "__main__":
    main()
